#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Fill Mask Caption脚本
使用GPT-4o-mini为mask图像生成caption，支持多线程处理和JSON schema约束
"""

import os
import json
import base64
import requests
import cv2

from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Dict, List, Any, Optional
import time
import logging
from datetime import datetime
from tqdm import tqdm

# 配置常量
API_KEY = "your_api_key_here"
BASE_URL = "https://openrouter.ai/api/v1"
JSON_PATH = "/path/to/masked.json"
PROMPT_PATH = "/path/to/patch_generate_prompt.txt"
LOG_DIR = "/path/to/logs"

# 线程和API配置
MAX_WORKERS = 3  # 最大并发线程数
MAX_RETRIES = 3  # 最大重试次数
RETRY_DELAY = 2  # 重试间隔(秒)
REQUEST_TIMEOUT = 60  # 请求超时时间(秒)

# JSON Schema定义
RESPONSE_SCHEMA = {
    "type": "object",
    "properties": {
        "num_blocks": {
            "type": "integer",
            "minimum": 1,
            "maximum": 100
        },
        "mask": {
            "type": "array",
            "items": {
                "type": "integer",
                "minimum": 1
            }
        },
        "mask_caption": {
            "type": "string",
            "minLength": 10,
            "maxLength": 1000
        }
    },
    "required": ["num_blocks", "mask", "mask_caption"],
    "additionalProperties": False
}


class MaskCaptionGenerator:
    """Mask Caption生成器"""

    def __init__(self):
        self.api_key = API_KEY
        self.base_url = BASE_URL
        self.json_path = JSON_PATH
        self.prompt_path = PROMPT_PATH

        # 设置日志
        self.logger = self._setup_logging()

        # 加载prompt
        self.prompt_template = self._load_prompt()

        # 统计信息
        self.total_items = 0
        self.processed_items = 0
        self.failed_items = 0

    def _setup_logging(self) -> logging.Logger:
        """设置日志系统"""
        logger = logging.getLogger('MaskCaptionGenerator')
        logger.setLevel(logging.INFO)

        # 清除现有handlers
        if logger.handlers:
            logger.handlers.clear()

        # 创建formatter
        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')

        # 控制台handler
        console_handler = logging.StreamHandler()
        console_handler.setLevel(logging.INFO)
        console_handler.setFormatter(formatter)
        logger.addHandler(console_handler)

        # 文件handler
        log_dir = LOG_DIR
        log_file = os.path.join(log_dir, f"fill_mask_caption_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log")
        file_handler = logging.FileHandler(log_file, encoding='utf-8')
        file_handler.setLevel(logging.INFO)
        file_handler.setFormatter(formatter)
        logger.addHandler(file_handler)

        return logger

    def _load_prompt(self) -> str:
        """加载prompt模板"""
        try:
            with open(self.prompt_path, 'r', encoding='utf-8') as f:
                prompt = f.read().strip()
            self.logger.info(f"Prompt模板加载成功: {self.prompt_path}")
            return prompt
        except Exception as e:
            self.logger.error(f"加载prompt模板失败: {e}")
            raise

    def _load_json_data(self) -> List[Dict[str, Any]]:
        """加载JSON数据"""
        try:
            with open(self.json_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
            self.logger.info(f"JSON数据加载成功: {self.json_path}, 共 {len(data)} 个项目")
            return data
        except Exception as e:
            self.logger.error(f"加载JSON数据失败: {e}")
            raise

    def _save_json_data(self, data: List[Dict[str, Any]], output_path: str):
        """保存JSON数据"""
        try:
            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(data, f, indent=4, ensure_ascii=False)
            self.logger.info(f"JSON数据保存成功: {output_path}")
        except Exception as e:
            self.logger.error(f"保存JSON数据失败: {e}")
            raise

    def _encode_image_to_base64(self, image_path: str) -> Optional[str]:
        """将图像编码为base64"""
        try:
            if not os.path.exists(image_path):
                self.logger.warning(f"图像文件不存在: {image_path}")
                return None

            # 读取图像
            image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
            if image is None:
                self.logger.warning(f"无法读取图像: {image_path}")
                return None

            # 编码为PNG格式的bytes
            success, buffer = cv2.imencode('.png', image)
            if not success:
                self.logger.warning(f"图像编码失败: {image_path}")
                return None

            # 转换为base64
            image_base64 = base64.b64encode(buffer.tobytes()).decode('utf-8')
            return image_base64

        except Exception as e:
            self.logger.error(f"图像编码错误 {image_path}: {e}")
            return None

    def _call_gpt_api(self, image_base64: str, item_id: str) -> Optional[Dict[str, Any]]:
        """调用GPT API"""
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }

        # 构建消息
        messages = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": self.prompt_template
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/png;base64,{image_base64}"
                        }
                    }
                ]
            }
        ]

        # 构建请求数据
        data = {
            "model": "openai/gpt-4o-mini",
            "messages": messages,
            "max_tokens": 800,
            "temperature": 0.1,
            "response_format": {
                "type": "json_schema",
                "json_schema": {
                    "name": "mask_analysis",
                    "schema": RESPONSE_SCHEMA,
                    "strict": True
                }
            }
        }

        # 发送请求
        for attempt in range(MAX_RETRIES):
            try:
                response = requests.post(
                    f"{self.base_url}/chat/completions",
                    headers=headers,
                    json=data,
                    timeout=REQUEST_TIMEOUT
                )

                if response.status_code == 200:
                    result = response.json()

                    if 'choices' in result and len(result['choices']) > 0:
                        content = result['choices'][0]['message']['content']

                        # 解析JSON响应
                        try:
                            parsed_result = json.loads(content)
                            self.logger.info(f"API调用成功 - {item_id}")
                            return parsed_result
                        except json.JSONDecodeError as e:
                            self.logger.warning(f"JSON解析失败 - {item_id}: {e}, 内容: {content}")
                            if attempt < MAX_RETRIES - 1:
                                time.sleep(RETRY_DELAY)
                                continue
                    else:
                        self.logger.warning(f"API响应格式错误 - {item_id}: {result}")
                else:
                    self.logger.warning(f"API请求失败 - {item_id}: {response.status_code}, {response.text}")

                # 如果不是最后一次尝试，等待后重试
                if attempt < MAX_RETRIES - 1:
                    time.sleep(RETRY_DELAY * (attempt + 1))  # 递增延迟

            except requests.exceptions.Timeout:
                self.logger.warning(f"API请求超时 - {item_id}, 尝试 {attempt + 1}/{MAX_RETRIES}")
                if attempt < MAX_RETRIES - 1:
                    time.sleep(RETRY_DELAY * (attempt + 1))
            except Exception as e:
                self.logger.error(f"API请求异常 - {item_id}: {e}")
                if attempt < MAX_RETRIES - 1:
                    time.sleep(RETRY_DELAY * (attempt + 1))

        self.logger.error(f"API调用最终失败 - {item_id}")
        return None

    def _process_single_item(self, item: Dict[str, Any], index: int) -> Dict[str, Any]:
        """处理单个项目"""
        item_id = f"{item.get('question_id', 'unknown')}_{index}"

        try:
            # 检查是否有mask_path
            mask_paths = item.get('mask_path', [])
            if not mask_paths:
                self.logger.warning(f"项目 {item_id} 没有mask_path")
                return item

            # 获取第一个mask路径
            mask_path = mask_paths[0] if isinstance(mask_paths, list) else mask_paths

            # 编码图像
            image_base64 = self._encode_image_to_base64(mask_path)
            if not image_base64:
                self.logger.warning(f"项目 {item_id} 图像编码失败")
                return item

            # 调用API
            api_result = self._call_gpt_api(image_base64, item_id)
            if api_result:
                # 更新项目数据
                item['mask'] = api_result.get('mask', [])
                item['mask_caption'] = api_result.get('mask_caption', '')

                # 可选：添加额外信息
                if 'num_blocks' in api_result:
                    item['num_blocks'] = api_result['num_blocks']

                self.processed_items += 1
                self.logger.info(f"项目 {item_id} 处理成功")
            else:
                self.failed_items += 1
                self.logger.error(f"项目 {item_id} 处理失败")

            return item

        except Exception as e:
            self.logger.error(f"处理项目 {item_id} 时出错: {e}")
            self.failed_items += 1
            return item

    def process_all_items(self):
        """处理所有项目"""
        self.logger.info("开始处理所有项目")

        # 加载数据
        data = self._load_json_data()
        self.total_items = len(data)

        if self.total_items == 0:
            self.logger.warning("没有找到需要处理的项目")
            return

        # 过滤出有mask_path的项目
        items_to_process = []
        for i, item in enumerate(data):
            mask_paths = item.get('mask_path', [])
            if mask_paths:
                items_to_process.append((item, i))

        self.logger.info(f"总项目数: {self.total_items}, 需要处理的项目数: {len(items_to_process)}")

        if not items_to_process:
            self.logger.warning("没有找到包含mask_path的项目")
            return

        # 多线程处理
        processed_data = data.copy()  # 复制原始数据

        with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
            # 提交任务
            future_to_index = {}
            for item, index in items_to_process:
                future = executor.submit(self._process_single_item, item, index)
                future_to_index[future] = index

            # 处理结果
            with tqdm(total=len(items_to_process), desc="处理项目") as pbar:
                for future in as_completed(future_to_index):
                    index = future_to_index[future]
                    try:
                        result = future.result()
                        processed_data[index] = result
                    except Exception as e:
                        self.logger.error(f"线程执行异常，索引 {index}: {e}")
                        self.failed_items += 1
                    finally:
                        pbar.update(1)

        # 保存结果
        output_path = self.json_path.rsplit('.', 1)[0] + '_ai.json'
        self._save_json_data(processed_data, output_path)

        # 输出统计信息
        self.logger.info("=" * 50)
        self.logger.info("处理完成统计:")
        self.logger.info(f"总项目数: {self.total_items}")
        self.logger.info(f"需要处理的项目数: {len(items_to_process)}")
        self.logger.info(f"成功处理: {self.processed_items}")
        self.logger.info(f"处理失败: {self.failed_items}")
        self.logger.info(f"成功率: {self.processed_items / len(items_to_process) * 100:.1f}%")
        self.logger.info(f"输出文件: {output_path}")
        self.logger.info("=" * 50)


def main():
    """主函数"""
    print("启动Fill Mask Caption处理程序")

    try:
        generator = MaskCaptionGenerator()
        generator.process_all_items()
        print("程序执行完成")
    except KeyboardInterrupt:
        print("\n程序被用户中断")
    except Exception as e:
        print(f"程序执行出错: {e}")
        logging.error(f"程序执行出错: {e}")


if __name__ == "__main__":
    main()
